#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0] == 'extract_edgels-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

from numpy import dot

from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors

from quaternion import Quat

import simplejson

import Image

set_printoptions(precision=3)

if __name__ == '__main__':

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0] == 'extract_edgels-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load image and initialize pic object

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])
    filename = fileroot+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], double(job_params["gaussian_smoothing_factor"]))


    ## Creates picture object

    if not 'projection_model' in job_params:
        model = 'perspective'
    else:
        model = job_params['projection_model']

    if model == 'perspective':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance_guess']
        p_point = array(job_params['principal_point_guess'])

        pic = PicturePinhole(imr, focal_distance, p_point[0], p_point[1])

    elif model == 'equirectangular':
        pic = PictureEqrec(imr)

    ## Don take last columns in consideration...
    #pic.mask[:,-3:] = 0 #remove last column
    ## load eqrec mask
    #imm = Image.open('eqrec-mask.png')
    #mask = np.array(imm.resize( (pic.Iwidth, pic.Iheight) ))
    #pic.mask = mask[:,:,0]


    ## Run the grid sweep, and generate the edgels list.
    ## Edgel extractor parameters
    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold'] # ** -1
    pic.p3 = job_params['angular_threshold'] ## this is for the orientation estimation
    pic.extract_edgels(gspc, glim)
    pic.calculate_extrinsic_edgels()
    ##
    #################################################################

    #################################################################
    ## Execute the optimizaiton
    xini = array(job_params['initial_guesses'])   ## Read initial points form job file
    #errs = [pic.old_find_orientation(xi) for xi in xini]
    errs = [pic.angle_find_orientation(xi) for xi in xini]
    xopt = Quat(errs[argmin([ee[1] for ee in errs])][0])
    ##
    #################################################################
    pic.orientation = xopt.canonical()
    print framenum, 'Ans:', xopt

    ## Make sure we have the right labels
    pic.calculate_MAP(pic.orientation.rot().T)

    ## Calculate rectified coordinates of observations, and create
    ## kd-trees to perform searches.
    gain = job_params['extraction_gain']
    pic.calculate_rectified_observations(gain)

    max_dist = job_params['extraction_distance']

    #################################################################
    ## Line extraction. For each direction, go through the list of
    ## edgels and look for the nearest neighbors to the left and
    ## right. If this neighborhood is nor registered already, put it
    ## on the lis tof neighbors.

    ## This can probably be better implemented, looks kludgy...

    nr={}
    nl={}

    for lab in range(3):
        for qi in range(0,len(pic.s_edgels[lab]),1):
            qs = pic.s_edgels[lab][qi]
            qo = pic.o_edgels[lab][qi]
            ## Fetch N nearest neighbours. Not approximate, minkowski norm (p=1)
            dnn,nn = pic.trees[lab].query(qs,20,0,1,max_dist)

            for dnq,ni in zip(dnn[1:],nn[1:]):
                if isinf(dnq):
                    break

                ns = pic.s_edgels[lab][ni]
                nd = ns - qs

                if np.abs(nd[0]) > np.abs(nd[1]):
                    continue

                if nd[1] > 0:
                    keyl = (lab,qi)
                    vl = (ni, dnq)
                    keyr = (lab,ni)
                    vr = (qi, dnq)
                else:
                    keyl = (lab,ni)
                    vl = (qi, dnq)
                    keyr = (lab,qi)
                    vr = (ni, dnq)

                if not keyl in nl and not keyr in nr:
                    nl[keyl]=vl
                    nr[keyr]=vr
                elif keyl in nl and not keyr in nr:
                    if nl[keyl][1]>dnq:
                        keyd=(lab,nl[keyl][0])
                        nr.pop(keyd)
                        nl.pop(keyl)
                        nl[keyl]=vl
                        nr[keyr]=vr
                elif not keyl in nl and keyr in nr:
                    if nr[keyr][1]>dnq:
                        keyd=(lab,nr[keyr][0])
                        nl.pop(keyd)
                        nr.pop(keyr)
                        nl[keyl]=vl
                        nr[keyr]=vr
                elif keyl in nl and keyr in nr :
                    if nl[keyl][1]>dnq and nr[keyr][1]>dnq:
                        keydl=(lab,nr[keyr][0])
                        keydr=(lab,nl[keyl][0])
                        nl.pop(keydl)
                        nr.pop(keydr)
                        nl[keyl]=vl
                        nr[keyr]=vr




    if PlotStuff:
        pic.calculate_edgel_normals()
        ## Plot estimated vanishing point directions
        figure(24, figsize=(12,9))

        title('Possible edge directions at multiple points')
        imshow(pic.frame/260., interpolation='nearest')
        pic.plot_vdirs(gca(), 25, pic.orientation.rot().T)
        axis('equal')
        savefig(fileroot+'/movie/vps-%04d.png'%framenum, dpi=150)

        # figure(25)
        # plot_aux.plot_spheres(xopt1.canonical.rot().T, pic.normals )
        figure(26)
        plot_aux.plot_spheres(xopt.canonical().rot().T, pic.normals )
        # figure(27)
        # pp = mgrid[0.:rezl1.shape[0]]/rezl1.shape[0]
        # plot(rezl1, pp)
        # plot(reztu, pp)
        # grid()

    arnl = array([ (l,k,nl[(l,k)][0]) for l,k in nl ], dtype=np.int32)
    arnr = array([ (l,k,nr[(l,k)][0]) for l,k in nr ], dtype=np.int32)

    outfile = open(fileroot+'/edgels/edgels-%04d.npz'%framenum, 'w')
    savez(outfile,
          xopt = pic.orientation.q,
          i_edgels0 = pic.i_edgels[0], 
          i_edgels1 = pic.i_edgels[1], 
          i_edgels2 = pic.i_edgels[2],
          o_edgels0 = pic.o_edgels[0], 
          s_edgels0 = pic.s_edgels[0], 
          descriptors0 = pic.descriptors[0], 
          o_edgels1 = pic.o_edgels[1], 
          s_edgels1 = pic.s_edgels[1], 
          descriptors1 = pic.descriptors[1], 
          o_edgels2 = pic.o_edgels[2], 
          s_edgels2 = pic.s_edgels[2], 
          descriptors2 = pic.descriptors[2], 
          arnl = arnl, 
          arnr = arnr
          )
    outfile.close()


    
